﻿using Microsoft.EntityFrameworkCore.DataEncryption;
using Microsoft.EntityFrameworkCore.DataEncryption.Providers;
using System;
using System.IO;
using System.Linq;
using System.Security.Cryptography;

namespace HiCore.WebApi.Example.DbContexts.Encryption
{
    /// <summary>
    /// Implements the Advanced Encryption Standard (AES) symmetric algorithm.
    /// </summary>
    public class AESProvider : IEncryptionProvider
    {
        /// <summary>
        /// AES block size constant.
        /// </summary>
        public const int AesBlockSize = 128;

        /// <summary>
        /// Initialization vector size constant.
        /// </summary>
        public const int InitializationVectorSize = 16;

        private readonly byte[] _key;
        private readonly CipherMode _mode;
        private readonly PaddingMode _padding;
        private readonly byte[] _iv;

        /// <summary>
        /// Creates a new <see cref="AesProvider"/> instance used to perform symmetric encryption and decryption on strings.
        /// </summary>
        /// <param name="key">AES key used for the symmetric encryption.</param>
        /// <param name="mode">Mode for operation used in the symmetric encryption.</param>
        /// <param name="padding">Padding mode used in the symmetric encryption.</param>
        public AESProvider(byte[] key, CipherMode mode = CipherMode.CBC, PaddingMode padding = PaddingMode.PKCS7)
        {
            _key = key;
            _mode = mode;
            _padding = padding;
        }

        /// <summary>
        /// Creates a new <see cref="AesProvider"/> instance used to perform symmetric encryption and decryption on strings.
        /// </summary>
        /// <param name="key">AES key used for the symmetric encryption.</param>
        /// <param name="initializationVector">AES Initialization Vector used for the symmetric encryption.</param>
        /// <param name="mode">Mode for operation used in the symmetric encryption.</param>
        /// <param name="padding">Padding mode used in the symmetric encryption.</param>
        public AESProvider(byte[] key, byte[] initializationVector, CipherMode mode = CipherMode.CBC, PaddingMode padding = PaddingMode.PKCS7) : this(key, mode, padding)
        {
            // Re-enabled to allow for a static IV.
            // This reduces security, but allows for encrypted values to be searched using LINQ.
            _iv = initializationVector;
        }

        /// <inheritdoc />
        public TStore Encrypt<TStore, TModel>(TModel dataToEncrypt, Func<TModel, byte[]> converter, Func<Stream, TStore> encoder)
        {
            if (converter is null)
            {
                throw new ArgumentNullException(nameof(converter));
            }

            if (encoder is null)
            {
                throw new ArgumentNullException(nameof(encoder));
            }

            byte[] data = converter(dataToEncrypt);
            if (data is null || data.Length == 0)
            {
                return default;
            }

            using var aes = CreateCryptographyProvider();
            using var memoryStream = new MemoryStream();

            byte[] initializationVector = _iv;
            if (initializationVector is null)
            {
                aes.GenerateIV();
                initializationVector = aes.IV;
                memoryStream.Write(initializationVector, 0, initializationVector.Length);
            }

            using var transform = aes.CreateEncryptor(_key, initializationVector);
            using var crypto = new CryptoStream(memoryStream, transform, CryptoStreamMode.Write);
            crypto.Write(data, 0, data.Length);
            crypto.FlushFinalBlock();

            memoryStream.Seek(0L, SeekOrigin.Begin);
            try
            {
                return encoder(memoryStream);
            }
            catch
            {
                return default;
            }
        }

        /// <inheritdoc />
        public TModel Decrypt<TStore, TModel>(TStore dataToDecrypt, Func<TStore, byte[]> decoder, Func<Stream, TModel> converter)
        {
            //data is null or ""?
            if (string.IsNullOrWhiteSpace(dataToDecrypt?.ToString()))
            {
                return default;
            }
            //data is base64 string?
            if (!StringHelper.IsBase64(dataToDecrypt.ToString()))
            {
                return default;
            }

            if (decoder is null)
            {
                throw new ArgumentNullException(nameof(decoder));
            }

            if (converter is null)
            {
                throw new ArgumentNullException(nameof(converter));
            }

            byte[] data = decoder(dataToDecrypt);
            if (data is null || data.Length == 0)
            {
                return default;
            }

            using var memoryStream = new MemoryStream(data);

            byte[] initializationVector = _iv;
            if (initializationVector is null)
            {
                initializationVector = new byte[InitializationVectorSize];
                memoryStream.Read(initializationVector, 0, initializationVector.Length);
            }

            using var aes = CreateCryptographyProvider();
            using var transform = aes.CreateDecryptor(_key, initializationVector);
            using var crypto = new CryptoStream(memoryStream, transform, CryptoStreamMode.Read);
            try
            {
                return converter(crypto);
            }
            catch
            {
                return default;
            }
        }
        /// <summary>
        /// Generates an AES cryptography provider.
        /// </summary>
        /// <returns></returns>
        private AesCryptoServiceProvider CreateCryptographyProvider()
        {
            return new AesCryptoServiceProvider
            {
                BlockSize = AesBlockSize,
                Mode = _mode,
                Padding = _padding,
                Key = _key,
                KeySize = _key.Length * 8
            };
        }
    }

    internal class StringHelper
    {
        public static bool IsBase64(string text)
        {
            byte[] bytes = null;
            if (string.IsNullOrEmpty(text))
                return false;
            else
            {
                if (text.Length % 4 != 0)
                    return false;
                if (text.Any(c => !base64CodeArray.Contains(c)))
                    return false;
            }
            try
            {
                bytes = Convert.FromBase64String(text);
                return true;
            }
            catch (FormatException)
            {
                return false;
            }
        }

        private static char[] base64CodeArray = new char[]
        {
            'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
            'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
            '0', '1', '2', '3', '4',  '5', '6', '7', '8', '9', '+', '/', '='
        };
    }
}
